import os
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(project_root)
import torch
import matplotlib
matplotlib.use('Agg')

import numpy as np
import matplotlib.pyplot as plt

from concurrent.futures import ThreadPoolExecutor
from PIL import Image

from utils import *

def resize_to_match(image, target_size):
    return np.array(Image.fromarray(image).resize(target_size, Image.BILINEAR))

def load_label(label_path):
    label = Image.open(label_path)
    label = np.array(label, dtype=np.float32) / 255.0
    return label

def overlay_input_label(input_tensor, label, output_path):
    try:
        input_image = np.max(input_tensor, axis=-1)
        input_image = (input_image - input_image.min()) / (input_image.max() - input_image.min())
        input_image = np.flipud(input_image)
        
        label_resized = resize_to_match(label, input_image.shape[::-1])

        overlay = 0.7 * input_image + 0.3 * label_resized

        fig, ax = plt.subplots(figsize=(10, 10))
        cax = ax.imshow(overlay, cmap='viridis', origin='lower')
        fig.colorbar(cax, ax=ax, label="Normalized Intensity")
        ax.set_title("Overlay of Input Tensor and Label")
        ax.axis("off")

        fig.savefig(output_path)
        plt.close(fig) 
    except Exception as e:
        print(f"Error: {e}")

def process_plot_check(tensor_folder, tensor_file, label_folder, output_folder):
    tensor_path = os.path.join(tensor_folder, tensor_file)
    label_path = os.path.join(label_folder, tensor_file.replace(".pt", ".png"))

    if not os.path.exists(label_path):
        print(f"Label not found for {tensor_file}, skipping.")
        return

    input_tensor = torch.load(tensor_path).numpy()
    label = load_label(label_path)

    output_path = os.path.join(output_folder, tensor_file.replace(".pt", ".png"))
    overlay_input_label(input_tensor, label, output_path)
    print(f"Overlay saved to {output_path}")

def main():
    tensor_folder = "./workspace/image/SAR/tensors"
    label_folder = "./workspace/image/SAR/label"
    output_folder = "./workspace/image/SAR/check"
    os.makedirs(output_folder, exist_ok=True)

    with ThreadPoolExecutor(max_workers=10) as executor:
        files = os.listdir(tensor_folder)
        files.sort(key=natural_key)
        for tensor_file in files:
            if tensor_file.endswith(".pt"):
                executor.submit(process_plot_check, tensor_folder, tensor_file, label_folder, output_folder)

if __name__ == "__main__":
    main()